Transfer the code mainly on host offloading from g3 to github:

• b/816371919, Misc fixes to maxdiffusion: 1. Shards the hidden state input to the model 2. Factor in different flash block sizes to estimate the final padding for q/k/v.
• b/817363551, Fix fused splash kernel configuration
• b/817366090, Add remat policy for hidden states and attention activations

Test on v5p-8:

• FULL Remat, Step time 36718.2 ms, Peak Memory 49.37GB
• HIDDEN_STATE_WITH_OFFLOAD Remat, Step time 32776.9 ms, Peak Memory 48.72 GB